'''
Author: zhanwei xu
Date: 2023-06-21 16:37:17
LastEditors: zhanwei xu
LastEditTime: 2023-06-21 23:18:22
Description: 

Copyright (c) 2023 by zhanwei xu, Tsinghua University, All Rights Reserved. 
'''

from .utils import *
from . import bingsearch
from fastapi.responses import StreamingResponse, JSONResponse
from datetime import datetime
import threading
import time
import os
import poe

current_time = datetime.now()
print(current_time)

def save_count(chatbot):
    while True:
        current_time = datetime.now()
        if current_time.hour == 1 and current_time.minute == 0: #检查时间是否凌晨1点
            with open('count.txt', 'a') as f: # 打开文件以追加的方式保存count
                f.write(str(current_time) + ': ' + str(chatbot.count) + '\n')  # 保存count
            chatbot.count = 0  # 重置count
            time.sleep(60)  # 休眠60秒，以避免重复保存
        else:
            time.sleep(10) # 休眠10秒，然后再次检查时间

class poebot():  
    def __init__(self) -> None:
        self.ip_traffic = {}
        self.ip_traffic_4 = {}
        self.api_key_index = 0
        self.clients = self.get_clients()
        self.api_key_index = 0
        self.count = 0
        self.count_thread = threading.Thread(target=save_count, args=(self,))
        self.count_thread.daemon = True # 设置为守护线程，这样在主程序退出时，线程也会自动退出
        self.count_thread.start() # 启动线程
    def get_clients(self):
        
        f = open('p_b_value.txt', 'r')
        lines = f.readlines()
        f.close()
        clients = []
        for line in lines:
            try:
                clients.append(poe.Client(line.strip()))
            except Exception as e:
                print(e)
                continue
        return clients
    async def poechat(self, request):
        client_ip = request.client.host
        
        json_data = await request.json()
        
        # Extract message from request
        message = json_data["message"]
        try:
            mode = json_data["mode"]
        except:
            mode = 'a2'
        # check if mode is empty
        if mode == '':
            mode = 'a2'
        if mode!= 'a2':
            key = json_data["key"]
            key_list = os.listdir("keys")
            # check if the key is valid
            if key not in key_list:
                # Check if IP address has exceeded traffic limit
                if check_traffic_poe(client_ip):
                    return StreamingResponse(generate_error_poe(), media_type='application/json')
                else:
                    flag = 0
            # check if the key is empty
            elif os.stat("keys/" + key).st_size == 0:
                
                # Check if IP address has exceeded traffic limit
                if check_traffic_poe(client_ip):
                    return StreamingResponse(generate_error_poe(), media_type='application/json')
                else:
                    flag = 0
            # Extract message from request
            message = json_data["message"]
            if flag:
                key_path = "keys/" + key
                with open(key_path, 'r') as f:
                    key_conunt = f.read()
                    print(key_conunt)
                    key_conunt = int(key_conunt)
                    key_count = key_conunt -1
                with open(key_path, 'w') as f:
                    f.write(str(key_count))
                if key_count <= 0:
                    os.remove(key_path)
        prompt = [message[-1]]        
        message = message[-1]["content"]
        # print(prompt)
        history_token_num = num_tokens_from_message(prompt)
        if history_token_num > 100096:
            return StreamingResponse(generate_error_max(), media_type="text/plain")

        client = self.clients[self.api_key_index]
        self.api_key_index = (self.api_key_index + 1) % len(self.clients)
        
        return StreamingResponse(generate_text_poe(client,message,mode=mode), media_type="text/plain" )

